
import numpy as np

from view.endSrc.ProbClusters import ProbClusters
from view.endSrc.MyTools import MyTools

class GaussianClusters(ProbClusters):

    def __new__(cls, meanList, varMatList, paiList, probMatrix=None, labels=None, weight=float('inf')):
        '''
            From the MixedGaussianDatasetGenerator, this factory
            is used to create the objects of GaussianClusters

            WARNING: can not remove this function, or it will happen
            TypeError: __new__() got multiple values for argument 'probMatrix'

        :param probMatrix:  np.ndarray with shape (n,k), n-numOfRecords, k-numOfClusters
        :param labels:      the list of cluster ids for each record in the dataset

        :param meanList:    the list of mean scalar/vector for a d-dimensinal Gaussian mixture
                            the length of list is the number of components of the Gaussian mixture

        :param varMatList:  the list of variance scalar/matrix for a d-dimensinal Gaussian mixture
                            the length of list is the number of components of the Gaussian mixture

        :param paiList:     the list of ratios(scalar) for a d-dimensional Gaussian mixture
        :return: an object of the class GaussianClusters
        '''

        self = ProbClusters.__new__(cls, probMatrix, labels, weight)

        return self

    def __init__(self, meanList, varMatList, paiList, probMatrix=None, labels=None, weight=float('inf')):
        '''
                This method is called after GaussianClusters.__new__,
                so that we can check if the nClusters of probMatrix and paiList are same.
        :param meanList:
        :param varMatList:
        :param paiList:
        :param probMatrix:
        :param labels:
        '''
        ProbClusters.__init__(self, probMatrix, labels, weight)

        ###################################################################
        # check parameters of a Gaussian mixture
        assert isinstance(meanList, list), 'meanList must be a list'
        assert isinstance(varMatList, list), 'varMatList must be a list'

        paiList = np.asarray(paiList)
        paiList = paiList.astype(np.float64)
        assert len(paiList.shape) == 1, 'paiList must be a list or np.ndarray with shape like (k, )'

        s = sum(paiList)
        assert s == 1.0 or (abs(1.0-s) < 1e-10), 'sum of paiList must be 1.0, but it is ' + str(s)

        nClusters = len(meanList)
        assert nClusters == len(varMatList), 'all lists must have same number of components(clusters)'
        assert nClusters == len(paiList), 'all lists must have same number of components(clusters)'

        shp = np.shape(meanList[0])
        shpLen = len(shp)
        assert  shpLen <= 1, 'GaussianClusters.init. all members in meanList must be scalars or vectors'
        shp2 = () if shpLen == 0 else (shp[0], shp[0])

        MyTools.checkMemberShape(meanList, shp, 'meanList')
        MyTools.checkMemberShape(varMatList, shp2, 'varMatList')
        MyTools.checkMemberShape(paiList, (), 'varMatList')

        ###################################################################
        # add checked specific properties of GaussianClusters
        self.meanList = meanList
        self.varMatList = varMatList
        self.paiList = paiList.tolist()   # to list of float, so that we can transform it to JSON list

        # check nClusters, since ProbClusters can extract numOfClusters from ProbMatrix/labels without
        # we do it again in GaussianClusters.__new__
        assert self.getNumOfClusters() == len(paiList), \
            'numOfClusters of probMatrix/labels must be same as that of paiList/meanList/varMatList'


    @staticmethod
    def fromJSONDict(mixedParams, labels):
        '''
            From DB tDataset, this factory is used to
            create objects of GaussianClusters

        :param self:
        :param mixedParams: JSON Dict, i.e., this is no np.ndarray (all transformed to lists)
        :param labels: list<int>
        :return: a new created GaussianClusters
        '''
        assert MyTools.isJsonDumpable(mixedParams), 'GaussianClusters.fromJSONDict..mixedParams must be JSON dumpable'

        nClusters = len(np.unique(labels))

        dimVec = len(mixedParams['meanList'][0]) # get length of the mean vector from 1st element
        meanList = []
        for s in mixedParams['meanList']:
            if isinstance(s, float) or \
                isinstance(s, np.float64) or \
                isinstance(s, int) or \
                isinstance(s, np.int64):
                meanList.append(s)              # scalar case
                continue
            meanList.append(np.asarray(s))      # list case, from list to np.ndarray

        varMatList = []
        for s in mixedParams['varMatList']:
            if isinstance(s, float) or \
                isinstance(s, np.float64) or \
                isinstance(s, int) or \
                isinstance(s, np.int64):
                varMatList.append(s)        # scalar case
                continue
            s = np.asarray(s)               # list case: from list to np.ndarray
            s.reshape((dimVec, dimVec))     # from list<list> to np.ndarray 2D square matrix
            varMatList.append(s)

        paiList = mixedParams['paiList']   # list of float, will be checked in __init__

        return GaussianClusters(meanList, varMatList, paiList, labels=labels)

    def dumpParamsToJSONDict(self):
        '''
            This method is used to get the JSON dict representation of  self.meanList,varMatList,paiList,
            so that user can write the JSON dict representation to DB
        :return: JSON dict represenation for self.meanList,varMatList, paiList
        '''

        mp = {}

        meanList = []
        for s in self.meanList:
            if isinstance(s, float) or \
                isinstance(s, np.float64) or \
                isinstance(s, int) or \
                isinstance(s, np.int64):
                meanList.append(s)              # scalar case
            else:
                meanList.append(s.tolist())     # numpy case:  tolist
        mp['meanList'] = meanList

        varMatList = []
        for s in self.varMatList:
            if isinstance(s, float) or \
                isinstance(s, np.float64) or \
                isinstance(s, int) or \
                isinstance(s, np.int64):
                varMatList.append(s)            # scalar case
            else:
                varMatList.append(s.tolist())  # numpy case: tolist
        mp['varMatList'] = varMatList

        mp['paiList'] = self.paiList       # it is list of float, which is promised in __init__



        return mp


    def __str__(self):
        '''
            for simple test
        :return: a string representation of this object
        '''

        tmp = '---------- GaussianClusters ----------------\n'
        tmp += ProbClusters.__str__(self)

        tmp += 'dataDim: ' + str(self.getDataDimension()) + '\n'
        tmp += 'paiList: ' + str(self.paiList) + '\n'

        tmp += 'meanList: \n'
        for m in self.meanList:
            tmp += '\t' + str(m) + '\n'

        tmp += 'varMatList: \n'
        for v in self.varMatList:
            tmp += '\t' + str(v) + '\n'

        return tmp

    def getDataDimension(self):
        '''
            Must provide an implementation of the virtual method of superclass ProbClusters
        :return: the dimension of the mean parameter
                 0 if the mean is a scalar, >=1 if the mean is a np.ndarray vector
        '''
        a = self.meanList[0]
        if isinstance(a, np.ndarray):
            return len(a)
        # must be scalar, had been promised in GaussianClusters.__new__
        return 0



    @staticmethod
    def testInit():

        probMat = np.asarray([[0.2, 0.3, 0.5],
                              [0.1, 0.6, 0.3],
                              [0.5, 0.3, 0.1],
                              [0.6, 0.2, 0.2]])
        # labels: [[2, 3], [1], [0]]
        nClusters = probMat.shape[1]
        meanList= [ np.asarray([1.1, 2.2]),
                    np.asarray([3.4, 6.4]),
                    np.asarray([2.3, 5.6])
                  ] # nClusters = 3
        varMatList = [
            np.asarray([[2.1, 0.5],
                        [11.2, 9.2]]),
            np.asarray([[2, 3.1],
                        [4, 5.1]]),
            np.asarray([[0.1, 0],
                        [2, 1]]),
        ] # nClusters =3
        paiList = [2.1, 4.3, 89.1]
        s = sum(paiList)
        paiList = [x/s for x in paiList]

        gc = GaussianClusters(meanList, varMatList, paiList, probMatrix=probMat)
        print(gc)

    @staticmethod
    def testJSONDict():
        probMat = np.asarray([[0.2, 0.3, 0.5],
                              [0.1, 0.6, 0.3],
                              [0.5, 0.3, 0.1],
                              [0.6, 0.2, 0.2]])
        # labels: [[2, 3], [1], [0]]
        nClusters = probMat.shape[1]
        meanList= [ np.asarray([1.1, 2.2]),
                    np.asarray([3.4, 6.4]),
                    np.asarray([2.3, 5.6])
                  ] # nClusters = 3
        varMatList = [
            np.asarray([[2.1, 0.5],
                        [11.2, 9.2]]),
            np.asarray([[2, 3.1],
                        [4, 5.1]]),
            np.asarray([[0.1, 0],
                        [2, 1]]),
        ] # nClusters =3
        paiList = [2.1, 4.3, 89.1]
        s = sum(paiList)
        paiList = [x/s for x in paiList]

        gc = GaussianClusters(meanList, varMatList, paiList, probMatrix=probMat)
        print('.......before dump....', gc)

        labels = gc.getLabels()

        ##########################
        mpDic = gc.dumpParamsToJSONDict()

        gc2 = GaussianClusters.fromJSONDict(mpDic, labels=labels)

        print('........reload from JSON...', gc2)


if __name__ == '__main__':

    # GaussianClusters.testInit()
    GaussianClusters.testJSONDict()